from models.resnet import resnet18
from models.resnet_prune import resnet18_prune
from models.shufflenet import ShuffleNetV2
from models.tiny_mobilenet_v2 import TinyMobileNetV2
from models.shufflenet_prune import ShuffleNetV2Prune

MODEL = {'resnet18': resnet18,
         'resnet18_prune': resnet18_prune,
         'shufflenet': ShuffleNetV2,
         'mobilenet': TinyMobileNetV2,
         'shufflenet_prune': ShuffleNetV2Prune}